// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <condition_variable>
#include <mutex>

#include "melon/exception_wrapper.h"
#include "melon/executor.h"

namespace kumo::pollux::dwio::common {
    class ExecutorBarrier : public melon::Executor {
    public:
        explicit ExecutorBarrier(melon::Executor::KeepAlive<> executor)
            : executor_{std::move(executor)}, count_{0} {
        }

        // Constructor version that holds ownership over the executor (holds a
        // shared_ptr copy).
        explicit ExecutorBarrier(std::shared_ptr<melon::Executor> executor)
            : owned_{std::move(executor)},
              executor_{melon::getKeepAliveToken(*owned_)},
              count_{0} {
        }

        ~ExecutorBarrier() override {
            // If this object gets destroyed while there are still tasks pending, those
            // tasks will try to access invalid memory addresses in the current object.
            std::unique_lock lock{mutex_};
            cv_.wait(lock, [&]() { return count_ == 0; });
            // We won't throw from the destructor so we don't check for exceptions
            // Also, I don't need to clear the exception because this is the destructor.
        }

        /// Enqueue a function to be executed by this executor. This and all
        /// variants must be threadsafe.
        void add(melon::Func) override;

        /// Enqueue a function with a given priority, where 0 is the medium priority
        /// This is up to the implementation to enforce
        void addWithPriority(melon::Func, int8_t priority) override;

        uint8_t getNumPriorities() const override {
            return executor_->getNumPriorities();
        }

        void waitAll() {
            std::unique_lock lock{mutex_};
            cv_.wait(lock, [&]() { return count_ == 0; });
            if (exception_.has_exception_ptr()) {
                melon::exception_wrapper ew;
                // Clear the exception for the next time
                std::swap(ew, exception_);
                ew.throw_exception();
            }
        }

    private:
        auto wrapMethod(melon::Func f);

        std::shared_ptr<melon::Executor> owned_;
        melon::Executor::KeepAlive<> executor_;
        size_t count_;
        std::mutex mutex_;
        std::condition_variable cv_;
        melon::exception_wrapper exception_;
    };
} // namespace kumo::pollux::dwio::common
